import csv
import argparse
from typing import List, Dict, Tuple, Optional

class NetworkTopologyGenerator:
    """生成NVIDIA风格网络拓扑的CSV配置文件"""
    
    def __init__(self, server_count: int, 
                 server_template: str = "gpu{id}",
                 leaf_template: str = "leaf{id}",
                 spine_template: str = "spine{id}",
                 core_template: str = "core{id}",
                 server_ports: int = 8,
                 leaf_ports: int = 64,  # 修改为64端口
                 spine_ports: int = 64,  # 修改为64端口
                 core_ports: int = 64):  # 修改为64端口
        """
        初始化网络拓扑生成器
        
        参数:
            server_count: 服务器数量
            server_template: 服务器命名模板
            leaf_template: 叶子交换机命名模板
            spine_template: 脊交换机命名模板
            core_template: 核心交换机命名模板
            server_ports: 每台服务器的端口数
            leaf_ports: 每台叶子交换机的端口数
            spine_ports: 每台脊交换机的端口数
            core_ports: 每台核心交换机的端口数
        """
        self.server_count = server_count
        self.server_template = server_template
        self.leaf_template = leaf_template
        self.spine_template = spine_template
        self.core_template = core_template
        self.server_ports = server_ports
        self.leaf_ports = leaf_ports
        self.spine_ports = spine_ports
        self.core_ports = core_ports
        
        # 确保叶子交换机数量足够
        self.leaf_count = max(self.server_ports, self._calculate_leaf_count())
        self.spine_count = self._calculate_spine_count()
        self.core_count = self._calculate_core_count()
        
        # 存储连接信息
        self.connections: List[Dict[str, str]] = []
    
    def _calculate_leaf_count(self) -> int:
        """计算所需的叶子交换机数量"""
        # 每台叶子交换机的下联端口数 = 总端口数的一半 (64端口的交换机下联32个服务器)
        downlink_ports_per_leaf = self.leaf_ports // 2
        # 计算需要的叶子交换机数量
        return (self.server_count + downlink_ports_per_leaf - 1) // downlink_ports_per_leaf
    
    def _calculate_spine_count(self) -> int:
        """计算所需的脊交换机数量"""
        # 每台叶子交换机的上联端口数 = 总端口数的一半 (64端口的交换机上联32个脊)
        uplink_ports_per_leaf = self.leaf_ports // 2
        # 假设每台脊交换机连接所有叶子交换机
        return (self.leaf_count + uplink_ports_per_leaf - 1) // uplink_ports_per_leaf
    
    def _calculate_core_count(self) -> int:
        """计算所需的核心交换机数量"""
        # 假设每台核心交换机连接所有脊交换机
        return 1  # 简化为1台核心交换机
    
    def _get_device_name(self, template: str, id: int) -> str:
        """根据模板生成设备名称"""
        return template.format(id=id)
    
    def generate_server_to_leaf_connections(self) -> None:
        """生成服务器到叶子交换机的连接（每台服务器的8个端口连接到8个不同的叶子交换机）"""
        for server_id in range(1, self.server_count + 1):
            server_name = self._get_device_name(self.server_template, server_id)
            
            for port_idx in range(self.server_ports):
                # 确定连接的叶子交换机ID (循环使用叶子交换机)
                leaf_id = ((server_id - 1) + port_idx) % self.leaf_count + 1
                leaf_name = self._get_device_name(self.leaf_template, leaf_id)
                
                # 服务器端口从1开始
                server_port = port_idx + 1
                
                # 叶子交换机下联端口从1开始，每台服务器使用一个专用端口
                # 对于64端口交换机，下联端口范围是1-32
                leaf_port = (server_id - 1) % (self.leaf_ports // 2) + 1
                
                self.connections.append({
                    "adev": server_name,
                    "aport": f"eth{server_port}",
                    "zdev": leaf_name,
                    "zport": f"Ethernet{leaf_port}"
                })
    
    def generate_leaf_to_spine_connections(self) -> None:
        """生成叶子交换机到脊交换机的连接"""
        for leaf_id in range(1, self.leaf_count + 1):
            leaf_name = self._get_device_name(self.leaf_template, leaf_id)
            
            # 叶子交换机上联端口从总端口数的一半后开始 (对于64端口交换机，上联端口范围是33-64)
            start_port = (self.leaf_ports // 2) + 1
            
            for spine_id in range(1, self.spine_count + 1):
                spine_name = self._get_device_name(self.spine_template, spine_id)
                leaf_port = start_port + (spine_id - 1) % (self.leaf_ports // 2)
                
                # 脊交换机端口从1开始
                spine_port = (leaf_id - 1) % self.spine_ports + 1
                
                self.connections.append({
                    "adev": leaf_name,
                    "aport": f"Ethernet{leaf_port}",
                    "zdev": spine_name,
                    "zport": f"Ethernet{spine_port}"
                })
    
    def generate_spine_to_core_connections(self) -> None:
        """生成脊交换机到核心交换机的连接"""
        if self.core_count < 1:
            return
            
        for spine_id in range(1, self.spine_count + 1):
            spine_name = self._get_device_name(self.spine_template, spine_id)
            
            for core_id in range(1, self.core_count + 1):
                core_name = self._get_device_name(self.core_template, core_id)
                
                # 脊交换机连接核心的端口从总端口数的一半后开始 (对于64端口交换机，端口范围是33-64)
                spine_port = (self.spine_ports // 2) + (core_id - 1)
                
                # 核心交换机端口从1开始
                core_port = (spine_id - 1) % self.core_ports + 1
                
                self.connections.append({
                    "adev": spine_name,
                    "aport": f"Ethernet{spine_port}",
                    "zdev": core_name,
                    "zport": f"Ethernet{core_port}"
                })
    
    def generate_all_connections(self) -> None:
        """生成所有连接"""
        self.generate_server_to_leaf_connections()
        self.generate_leaf_to_spine_connections()
        self.generate_spine_to_core_connections()
    
    def write_to_csv(self, filename: str = "nvidia_topology.csv") -> None:
        """将连接信息写入CSV文件"""
        if not self.connections:
            self.generate_all_connections()
            
        with open(filename, 'w', newline='') as csvfile:
            fieldnames = ['adev', 'aport', 'zdev', 'zport']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            
            writer.writeheader()
            for connection in self.connections:
                writer.writerow(connection)
        
        print(f"已成功生成网络拓扑配置文件: {filename}")
        print(f"服务器数量: {self.server_count}")
        print(f"叶子交换机数量: {self.leaf_count}")
        print(f"脊交换机数量: {self.spine_count}")
        print(f"核心交换机数量: {self.core_count}")
        print(f"总连接数: {len(self.connections)}")

def main():
    """主函数，处理命令行参数并生成网络拓扑"""
    parser = argparse.ArgumentParser(description='生成NVIDIA风格网络拓扑的CSV配置文件')
    parser.add_argument('-s', '--servers', type=int, required=True, help='服务器数量')
    parser.add_argument('--server-template', type=str, default="gpu{id}", help='服务器命名模板，例如"gpu{id}"')
    parser.add_argument('--leaf-template', type=str, default="leaf{id}", help='叶子交换机命名模板，例如"leaf{id}"')
    parser.add_argument('--spine-template', type=str, default="spine{id}", help='脊交换机命名模板，例如"spine{id}"')
    parser.add_argument('--core-template', type=str, default="core{id}", help='核心交换机命名模板，例如"core{id}"')
    parser.add_argument('--server-ports', type=int, default=8, help='每台服务器的端口数')
    parser.add_argument('--leaf-ports', type=int, default=64, help='每台叶子交换机的端口数')
    parser.add_argument('--spine-ports', type=int, default=64, help='每台脊交换机的端口数')
    parser.add_argument('--core-ports', type=int, default=64, help='每台核心交换机的端口数')
    parser.add_argument('-o', '--output', type=str, default="nvidia_topology.csv", help='输出CSV文件路径')
    
    args = parser.parse_args()
    
    generator = NetworkTopologyGenerator(
        server_count=args.servers,
        server_template=args.server_template,
        leaf_template=args.leaf_template,
        spine_template=args.spine_template,
        core_template=args.core_template,
        server_ports=args.server_ports,
        leaf_ports=args.leaf_ports,
        spine_ports=args.spine_ports,
        core_ports=args.core_ports
    )
    
    generator.write_to_csv(args.output)

if __name__ == "__main__":
    main()    